#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import re
import os
import argparse
import random
from typing import List, Tuple, Optional, Dict, Any

class SwiftCodeReorder:
    def __init__(self):
        # 方法模式
        self.method_pattern = re.compile(
            r'^(?P<indent>\s*)'  # 缩进
            r'(?P<attributes>(@\w+\s+)*)'  # 属性注解
            r'(?P<access>public\s+|private\s+|internal\s+|fileprivate\s+)?'  # 访问控制
            r'(?P<override>override\s+)?'  # override 关键字
            r'(?P<static>static\s+|class\s+)?'  # static/class 关键字
            r'(?P<mutating>mutating\s+)?'  # mutating 关键字
            r'(?P<func>func\s+)'  # func 关键字
            r'(?P<name>\w+)'  # 方法名
            r'(?P<generic><[^>]*>)?'  # 泛型参数
            r'(?P<params>\([^{]*\))'  # 参数列表
            r'(?P<throws>\s+throws)?'  # throws
            r'(?P<rethrows>\s+rethrows)?'  # rethrows
            r'(?P<return>\s*->\s*[^{]+)?'  # 返回类型
            r'\s*\{'  # 开始大括号
        )
        
        # 属性模式
        self.property_pattern = re.compile(
            r'^(?P<indent>\s*)'  # 缩进
            r'(?P<attributes>(@\w+\s+)*)'  # 属性注解
            r'(?P<access>public\s+|private\s+|internal\s+|fileprivate\s+)?'  # 访问控制
            r'(?P<static>static\s+|class\s+)?'  # static/class 关键字
            r'(?P<lazy>lazy\s+)?'  # lazy 关键字
            r'(?P<mutating>mutating\s+)?'  # mutating 关键字
            r'(?P<let>let\s+|var\s+)'  # let/var 关键字
            r'(?P<name>\w+)'  # 属性名
            r'(?P<type>:[^=]+)?'  # 类型注解
            r'(?P<value>\s*=.*)?'  # 初始值
        )
        
        # 计算属性模式（带 get/set）
        self.computed_property_pattern = re.compile(
            r'^(?P<indent>\s*)'  # 缩进
            r'(?P<attributes>(@\w+\s+)*)'  # 属性注解
            r'(?P<access>public\s+|private\s+|internal\s+|fileprivate\s+)?'  # 访问控制
            r'(?P<static>static\s+|class\s+)?'  # static/class 关键字
            r'(?P<var>var\s+)'  # var 关键字
            r'(?P<name>\w+)'  # 属性名
            r'(?P<type>:[^\{]+)?'  # 类型注解
            r'\s*\{'  # 开始大括号
        )
        
        # 分类和类模式
        self.extension_pattern = re.compile(r'^\s*extension\s+(\w+(?:\.\w+)*)\s*\{')
        self.class_pattern = re.compile(r'^\s*(class|struct|enum|protocol)\s+(\w+)\s*\{')
        self.end_brace_pattern = re.compile(r'^\s*\}')

    def parse_swift_file(self, file_path: str) -> List[Dict[str, Any]]:
        """
        解析 Swift 文件，返回包含分类、属性和方法的结构
        """
        with open(file_path, 'r', encoding='utf-8') as file:
            lines = file.readlines()
        
        sections = []
        current_section = None
        current_properties = []
        current_methods = []
        brace_count = 0
        in_section = False
        section_start_line = 0
        
        i = 0
        while i < len(lines):
            line = lines[i].rstrip()
            
            # 检查是否开始新的分类或类
            extension_match = self.extension_pattern.match(line)
            class_match = self.class_pattern.match(line)
            
            if extension_match or class_match:
                # 保存前一个分类
                if current_section and (current_properties or current_methods):
                    sections.append({
                        'name': current_section,
                        'type': 'extension' if extension_match else 'class',
                        'properties': current_properties,
                        'methods': current_methods,
                        'start_line': section_start_line,
                        'end_line': i
                    })
                
                # 开始新的分类
                if extension_match:
                    current_section = f"extension {extension_match.group(1)}"
                else:
                    current_section = f"{class_match.group(1)} {class_match.group(2)}"
                
                current_properties = []
                current_methods = []
                brace_count = 1
                in_section = True
                section_start_line = i
                i += 1
                continue
            
            # 如果不在任何分类中，跳过
            if not in_section:
                i += 1
                continue
            
            # 检查属性定义（普通属性）
            property_match = self.property_pattern.match(line)
            if property_match and brace_count == 1 and not line.strip().endswith('{'):
                # 普通属性（单行）
                property_content = [line]
                current_properties.append({
                    'type': 'property',
                    'name': property_match.group('name'),
                    'content': property_content,
                    'is_lazy': bool(property_match.group('lazy'))
                })
                i += 1
                continue
            
            # 检查计算属性（带 get/set）
            computed_property_match = self.computed_property_pattern.match(line)
            if computed_property_match and brace_count == 1:
                property_content = [line]
                property_brace_count = 1
                j = i + 1
                
                # 收集完整的计算属性实现
                while j < len(lines) and property_brace_count > 0:
                    next_line = lines[j].rstrip()
                    property_content.append(next_line)
                    
                    if '{' in next_line:
                        property_brace_count += next_line.count('{')
                    if '}' in next_line:
                        property_brace_count -= next_line.count('}')
                    
                    j += 1
                
                current_properties.append({
                    'type': 'computed_property',
                    'name': computed_property_match.group('name'),
                    'content': property_content,
                    'is_lazy': False
                })
                i = j
                continue
            
            # 检查方法定义
            method_match = self.method_pattern.match(line)
            if method_match and brace_count == 1:
                method_content = [line]
                method_brace_count = 1
                j = i + 1
                
                # 收集完整的方法实现
                while j < len(lines) and method_brace_count > 0:
                    next_line = lines[j].rstrip()
                    method_content.append(next_line)
                    
                    if '{' in next_line:
                        method_brace_count += next_line.count('{')
                    if '}' in next_line:
                        method_brace_count -= next_line.count('}')
                    
                    j += 1
                
                current_methods.append({
                    'type': 'method',
                    'name': method_match.group('name'),
                    'content': method_content
                })
                i = j
                continue
            
            # 计算大括号数量
            if '{' in line:
                brace_count += line.count('{')
            if '}' in line:
                brace_count -= line.count('}')
            
            # 分类结束
            if brace_count == 0 and in_section:
                if current_section and (current_properties or current_methods):
                    sections.append({
                        'name': current_section,
                        'type': 'extension' if extension_match else 'class',
                        'properties': current_properties,
                        'methods': current_methods,
                        'start_line': section_start_line,
                        'end_line': i
                    })
                in_section = False
                current_section = None
                current_properties = []
                current_methods = []
            
            i += 1
        
        return sections

    def shuffle_section_members(self, section: Dict[str, Any]) -> Dict[str, Any]:
        """
        打乱分类中的属性和方法顺序
        """
        # 合并属性和方法
        all_members = section['properties'] + section['methods']
        
        # 打乱顺序
        random.shuffle(all_members)
        
        # 分离回属性和方法
        new_properties = [member for member in all_members if member['type'] in ['property', 'computed_property']]
        new_methods = [member for member in all_members if member['type'] == 'method']
        
        return {
            'name': section['name'],
            'type': section['type'],
            'properties': new_properties,
            'methods': new_methods,
            'start_line': section['start_line'],
            'end_line': section['end_line']
        }

    def shuffle_only_methods(self, section: Dict[str, Any]) -> Dict[str, Any]:
        """
        只打乱方法顺序，保持属性顺序
        """
        shuffled_methods = section['methods'].copy()
        random.shuffle(shuffled_methods)
        
        return {
            'name': section['name'],
            'type': section['type'],
            'properties': section['properties'],
            'methods': shuffled_methods,
            'start_line': section['start_line'],
            'end_line': section['end_line']
        }

    def shuffle_only_properties(self, section: Dict[str, Any]) -> Dict[str, Any]:
        """
        只打乱属性顺序，保持方法顺序
        """
        shuffled_properties = section['properties'].copy()
        random.shuffle(shuffled_properties)
        
        return {
            'name': section['name'],
            'type': section['type'],
            'properties': shuffled_properties,
            'methods': section['methods'],
            'start_line': section['start_line'],
            'end_line': section['end_line']
        }

    def write_swift_file(self, file_path: str, sections: List[Dict[str, Any]], original_lines: List[str]):
        """
        将修改后的内容写回文件
        """
        # 构建新文件内容
        new_lines = []
        section_index = 0
        i = 0
        
        while i < len(original_lines):
            line = original_lines[i].rstrip()
            
            # 检查是否到达一个分类的开始
            if section_index < len(sections):
                current_section = sections[section_index]
                
                extension_match = self.extension_pattern.match(line)
                class_match = self.class_pattern.match(line)
                
                section_name_match = False
                if extension_match:
                    section_name_match = f"extension {extension_match.group(1)}" == current_section['name']
                elif class_match:
                    section_name_match = f"{class_match.group(1)} {class_match.group(2)}" == current_section['name']
                
                if section_name_match:
                    # 添加分类/类定义
                    new_lines.append(line)
                    i += 1
                    
                    # 添加该分类的所有成员（属性和方法）
                    for member in current_section['properties'] + current_section['methods']:
                        for member_line in member['content']:
                            new_lines.append(member_line.rstrip())
                    
                    # 跳过原始文件中的内容，直到分类结束
                    brace_count = 1
                    while i < len(original_lines) and brace_count > 0:
                        next_line = original_lines[i].rstrip()
                        if '{' in next_line:
                            brace_count += next_line.count('{')
                        if '}' in next_line:
                            brace_count -= next_line.count('}')
                        i += 1
                    
                    # 添加结束大括号
                    if i < len(original_lines):
                        new_lines.append(original_lines[i-1].rstrip())
                    
                    section_index += 1
                    continue
            
            # 不在分类中的内容直接添加
            if i < len(original_lines):
                new_lines.append(original_lines[i].rstrip())
            i += 1
        
        # 写入文件
        with open(file_path, 'w', encoding='utf-8') as file:
            file.write('\n'.join(new_lines) + '\n')

    def process_file(self, file_path: str, shuffle_methods: bool = True, 
                    shuffle_properties: bool = True, seed: Optional[int] = None):
        """
        处理 Swift 文件，打乱方法和属性顺序
        """
        if seed is not None:
            random.seed(seed)
        
        print(f"处理文件: {file_path}")
        
        # 读取原始文件内容
        with open(file_path, 'r', encoding='utf-8') as file:
            original_lines = [line.rstrip() for line in file]
        
        # 解析文件结构
        sections = self.parse_swift_file(file_path)
        
        if not sections:
            print("未找到可处理的分类或类")
            return
        
        print(f"找到 {len(sections)} 个分类/类")
        
        # 处理每个分类
        modified_sections = []
        for section in sections:
            print(f"\n处理 {section['name']}:")
            print(f"  - 属性: {len(section['properties'])} 个")
            print(f"  - 方法: {len(section['methods'])} 个")
            
            if shuffle_methods and shuffle_properties:
                # 打乱所有成员
                modified_section = self.shuffle_section_members(section)
                print("  ✓ 打乱所有属性和方法顺序")
            elif shuffle_methods:
                # 只打乱方法
                modified_section = self.shuffle_only_methods(section)
                print("  ✓ 打乱方法顺序")
            elif shuffle_properties:
                # 只打乱属性
                modified_section = self.shuffle_only_properties(section)
                print("  ✓ 打乱属性顺序")
            else:
                modified_section = section
                print("  ⚠ 未进行任何打乱操作")
            
            modified_sections.append(modified_section)
        
        # 写回文件
        backup_path = file_path + '.bak'
        with open(backup_path, 'w', encoding='utf-8') as backup_file:
            backup_file.write('\n'.join(original_lines) + '\n')
        
        self.write_swift_file(file_path, modified_sections, original_lines)
        print(f"\n✓ 完成文件处理: {file_path}")
        print(f"✓ 原始文件已备份为: {backup_path}")

def main():
    parser = argparse.ArgumentParser(description='Swift 代码顺序打乱工具')
    parser.add_argument('file', help='要处理的 Swift 文件路径')
    parser.add_argument('--methods-only', action='store_true',
                       help='只打乱方法顺序')
    parser.add_argument('--properties-only', action='store_true',
                       help='只打乱属性顺序')
    parser.add_argument('--seed', type=int, help='随机种子，用于可重复的结果')
    
    args = parser.parse_args()
    
    if not os.path.exists(args.file):
        print(f"错误: 文件 '{args.file}' 不存在")
        return
    
    # 确定打乱模式
    shuffle_methods = not args.properties_only or args.methods_only
    shuffle_properties = not args.methods_only or args.properties_only
    
    if args.methods_only and args.properties_only:
        shuffle_methods = True
        shuffle_properties = True
    
    reorder = SwiftCodeReorder()
    reorder.process_file(
        args.file, 
        shuffle_methods=shuffle_methods,
        shuffle_properties=shuffle_properties,
        seed=args.seed
    )

if __name__ == "__main__":
    main()